// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/pipeline/threadpool_op.h>

#include <tbir/runtime/logging.h>
#include <tbir/runtime/threadpool/lock_based_thread_pool.h>
#include <tbir/runtime/threadpool/lock_free_thread_pool.h>

namespace tbir::runtime {

    void ThreadPoolOp::AtForkBefore() {
        pool_ = nullptr;
    }

    void ThreadPoolOp::AtForkAfterInParentOrChild() {
        this->Init();
    }

    void ThreadPoolOp::Init() {
        lock_free_ = GetAttr<bool>("lock_free");
        thread_nums_ = GetAttr<int32_t>("thread_nums");
        thread_name_ = GetAttr<Unicode>("thread_name");
        auto thread_name = thread_name_.encode();
        if (lock_free_) {
            pool_ = std::make_shared<internal::SPSCLockFreeThreadPool>(
                    thread_nums_, std::string(thread_name.data(), thread_name.size()));
        } else {
            pool_ = std::make_shared<internal::LockBasedThreadPool>(
                    thread_nums_, std::string(thread_name.data(), thread_name.size()));
        }
    }

    RTValue ThreadPoolOp::Process(PyArgs inputs) const {
        MXTHROW << "ThreadPoolOp can not be call directly!!!";
        return None;
    }

    TBIR_REGISTER_NATIVE_OP(ThreadPoolOp);

}  // namespace tbir::runtime